import argparse
import os
import sys
from pathlib import Path

import numpy as np
import torch
import torch.backends.cudnn as cudnn
from scipy import stats
import time
import urllib.request
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.common import DetectMultiBackend
from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device, time_sync

from myUtils.yolov5s import YOLOv5s
from myUtils.queueinfodete import QueueInformationDetection
from myUtils.portcommunication import PortCommunication
from myUtils.bc25 import BC25, BC25LineDetectorData
from myUtils.resultingdata import ResultingData

# --------------多线程工具--------------------
import threading
from queue import Queue
lock = threading.Lock()
q_data = Queue()
# --------------------------------------------

# ---------------源选择------------------------
# source = '0'
source = 'http://10.2.131.220:5555/video_feed0'
# source = r'./data/videos'
# source = ROOT / 'data/images'
# ---------------------------------------------
FRAME_COUNT = 100


def send_data(data1, data2):
    # lock.acquire()
    bc25_device.aliyun_mqtt_pub_data2(data1)

    bc25_device.aliyun_mqtt_pub_data3(data2)
    # lock.release()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')
    parser.add_argument('--source', type=str, default=source, help='file/dir/URL/glob, 0 for webcam')
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='inference size h,w')
    parser.add_argument('--conf-thres', type=float, default=0.55, help='confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='show results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--visualize', action='store_true', help='visualize features')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--line-thickness', default=1, type=int, help='bounding box thickness (pixels)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
    parser.add_argument('--queue-info', default=False)
    parser.add_argument('--bc25-send-info', default=False)
    parser.add_argument('--bc25-port', type=str, default='/dev/ttyUSB0')
    opt = parser.parse_args()
    print(opt)

    # 输入源解析-----------------------------------------------------
    source = str(source)
    save_img = not opt.nosave and not source.endswith('.txt')

    is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
    webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)

    if is_url and is_file:
        source = check_file(source)  # download

    # 储存目录的创建-------------------------------------------------
    save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)  # increment run
    (save_dir / 'labels' if opt.save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

    # 加载YOLO模型--------------------------------------------------
    detect_model = YOLOv5s(**vars(opt))

    # 队伍检测初始化---------------------------------------------------
    if opt.queue_info:
        queue_info_detection = QueueInformationDetection()

    # 结果储存函数
    result_data = ResultingData()

    # bc25初始化----------------------------------------------------
    if opt.bc25_send_info:
        # 串口初始化
        port = PortCommunication(opt.bc25_port, 9600, 2)  # 端口和波特率
        if port.port_is_open:
            # bc25初始化
            bc25_device = BC25LineDetectorData(port)

            # MQTT初始化
            bc25_device.mqtt_init()
            bc25_device.aliyun_mqtt_pub_data(50)  # 发送一个测试数据50
            print('------------BC25 Init End------------')


        else:
            print('Unable to start cloud service!')

    # 对不同的输入源采用不同的数据加载方式
    if webcam:
        view_img = check_imshow()
        cudnn.benchmark = True  # set True to speed up constant image size inference

        dataset = LoadStreams(source, img_size=opt.imgsz, stride=detect_model.stride, auto=detect_model.pt)
        bs = len(dataset)  # batch_size
    else:
        dataset = LoadImages(source, img_size=opt.imgsz, stride=detect_model.stride, auto=detect_model.pt)
        bs = 1  # batch_size
    vid_path, vid_writer = [None] * bs, [None] * bs

    # Run inference
    detect_model.model.warmup(imgsz=(1 if detect_model.pt else bs, 3, *opt.imgsz))  # warmup

    # 计算帧-----------------------
    fps_time = 0
    fps = 0
    # 数据发送---------------------
    frame_count = 0
    people_in_line_data_list = []

    # 开始正式预测---------------------------------------------------------
    for path, im, im0s, vid_cap, s in dataset:
        # detected 包含了本章图片所有的预测信息，用一个列表将所有信息包裹
        detected, im = detect_model.detect(path, im, save_dir)
        key = cv2.waitKey(1)  # 停留25ms，当为0的时候则堵塞在第一帧不会继续下去
        if key == ord('q'):  # 当键入空格或者q时，则退出while循环
            break

        # 遍历所有的预测框
        for i, det in enumerate(detected):  # per image， i代表第i张图片
            if webcam:  # batch_size >= 1
                p, im0, frame = path[i], im0s[i].copy(), dataset.count
                s += f'{i}: '
            else:
                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path
            save_path = str(save_dir / p.name)  # im.jpg
            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
            s += '%gx%g ' % im.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            imc = im0.copy() if opt.save_crop else im0  # for save_crop

            detect_model.annotator_init(im0, opt.line_thickness)

            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results 计算各个类别的目标个数
                number_of_person = 0
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    if c == 0:
                        number_of_person = n
                    s += f"{n} {detect_model.names[int(c)]}{'s' * (n > 1)}, "  # add to string

                if opt.queue_info:
                    queue_info_detection.get_queue_info(im0, number_of_person)

                # Write results， 循环各个检测到的目标
                for *xyxy, conf, cls in reversed(det):
                    if opt.queue_info:
                        queue_info_detection.get_persons_position_list(xyxy)

                    if save_img or opt.save_crop or opt.view_img:  # Add bbox to image
                        c = int(cls)  # integer class
                        label = None if opt.hide_labels else (opt.names[c] if opt.hide_conf else f'{detect_model.names[c]} {conf:.2f}')
                        detect_model.annotator_plotbox(cls, opt.hide_labels, opt.hide_conf, conf, xyxy)
                    if opt.save_crop:
                        save_one_box(xyxy, imc, file=save_dir / 'crops' / detect_model.names[c] / f'{p.stem}.jpg', BGR=True)

                if opt.queue_info:
                    queue_info_detection.get_persons_number()
                    queue_info_detection.get_speed_of_the_queue(True, fps)

                    # 结果获取方式一：通过主函数中创建的类(存在问题)
                    # result_data.get_people_in_line(queue_info_detection.people_in_line)
                    # result_data.get_line_speed(queue_info_detection.line_speed)
                    # 结果获取方式二：通过queue_info_detection中创建的类
                    result_data = queue_info_detection.return_result()

                    print('people: ' + str(result_data.people_in_line))
                    print('speed: ' + str(result_data.line_speed))

                    frame_count += 1
                    people_in_line_data_list.append(int(result_data.people_in_line))

                    if opt.bc25_send_info and port.port_is_open and frame_count == FRAME_COUNT:
                        data1 = np.array(people_in_line_data_list)
                        data1 = stats.mode(data1)[0][0]
                        data2 = result_data.line_speed
                        # q_data.put([result_data.people_in_line, result_data.line_speed])
                        send_thread = threading.Thread(target=send_data,
                                                       args=(data1, data2),
                                                       name='sent_t')
                        send_thread.start()
                        frame_count = 0
                        people_in_line_data_list.clear()

            fps = 1.0 / (time.time() - fps_time)

            blk = np.zeros(im0.shape, np.uint8)
            cv2.rectangle(blk, (0, 0), (180, 65), (255, 255, 255), thickness=-1)  # 注意在 blk的基础上进行绘制；
            im0 = cv2.addWeighted(im0, 1.0, blk, 0.5, 1)

            im0 = cv2.putText(im0, 'FPS: %f' % fps,
                              (20, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
            fps_time = time.time()
            im0 = cv2.putText(im0, 'People: %s' % str(result_data.people_in_line),
                              (20, 35), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
            im0 = cv2.putText(im0, 'Speed: %s' % str(result_data.line_speed),
                              (20, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

            # Stream results
            # im0 = annotator.result()
            if opt.view_img:
                cv2.imshow(str(p), im0)
                cv2.waitKey(1)  # 1 millisecond

            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'image':
                    cv2.imwrite(save_path, im0)
                else:  # 'video' or 'stream'
                    if vid_path[i] != save_path:  # new video
                        vid_path[i] = save_path
                        if isinstance(vid_writer[i], cv2.VideoWriter):
                            vid_writer[i].release()  # release previous video writer
                        if vid_cap:  # video
                            fps = vid_cap.get(cv2.CAP_PROP_FPS)
                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        else:  # stream
                            fps, w, h = 20, im0.shape[1], im0.shape[0]
                        save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                    vid_writer[i].write(im0)


        # Print results
    # t = tuple(x / detect_model.seen * 1E3 for x in detect_model.dt)  # speeds per image
    # LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
    if opt.save_txt or save_img:
        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if opt.save_txt else ''
        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
    if opt.update:
        strip_optimizer(opt.weights)  # update model (to fix SourceChangeWarning)
    if opt.bc25_send_info and port.port_is_open:
        bc25_device.colse_mqtt()
